
#include "lerp_tiling.h"
#include "register/op_def_registry.h"
#include "tiling_common.h"
#include <cstdint>
namespace optiling {

static ge::graphStatus TilingFunc(gert::TilingContext *context) {

  LerpTilingData tiling;
  auto startShape = context->GetInputShape(0)->GetStorageShape();
  auto endShape = context->GetInputShape(1)->GetStorageShape();
  auto weightShape = context->GetInputShape(2)->GetStorageShape();
  // 判断shape情况
  auto size = 1;
  bool start_end_need_broadcast = false;
  auto maxDimNum =
      std::max(std::max(startShape.GetDimNum(), endShape.GetDimNum()),
               weightShape.GetDimNum());
  std::vector<int64_t> broadcastShape(MAX_SHAPE_DIM);
  std::vector<int64_t> startShapeData(MAX_SHAPE_DIM);
  std::vector<int64_t> endShapeData(MAX_SHAPE_DIM);
  std::vector<int64_t> weightShapeData(MAX_SHAPE_DIM);
  broadcastShape[0] = maxDimNum;
  broadcastShape[1] = 0;
  startShapeData[0] = maxDimNum;
  endShapeData[0] = maxDimNum;
  weightShapeData[0] = maxDimNum;
  for (auto i = 1; i <= maxDimNum; ++i) {
    auto index = -i;
    auto start_dim = getDimDefault1<int64_t>(startShape, index);
    auto end_dim = getDimDefault1<int64_t>(endShape, index);
    auto weight_dim = getDimDefault1<int64_t>(weightShape, index);
    auto ret = getBroadcastDim<int64_t>(start_dim, end_dim);
    if (!std::get<0>(ret)) {
      std::cerr << "stard 与 end 不能broadcast! index = " << index
                << ", start_dim = " << start_dim << ", end_dim = " << end_dim
                << std::endl;
    }

    ret = getBroadcastDim<uint32_t>(std::get<1>(ret), weight_dim);
    if (!std::get<0>(ret)) {
      std::cerr << "stard 与 end 不能broadcast! index = " << index
                << ", start_dim = " << start_dim << ", end_dim = " << end_dim
                << std::endl;
    }
    auto broadcastDim = std::get<1>(ret);
    size *= broadcastDim;
    auto shapeDataIndex = index + maxDimNum + 2;
    broadcastShape[shapeDataIndex] = broadcastDim;
    startShapeData[shapeDataIndex] = start_dim;
    endShapeData[shapeDataIndex] = end_dim;
    weightShapeData[shapeDataIndex] = weight_dim;
    startShapeData[1] = startShapeData[1] || broadcastDim != start_dim;
    endShapeData[1] = endShapeData[1] || broadcastDim != end_dim;
    weightShapeData[1] = weightShapeData[1] || broadcastDim != weight_dim;
  }
  tiling.set_size(size);
  tiling.set_broadcastShape(broadcastShape.data());
  tiling.set_startShape(startShapeData.data());
  tiling.set_endShape(endShapeData.data());
  tiling.set_weightShape(weightShapeData.data());
  print_shape_vector("broadcastShape", broadcastShape);
  print_shape_vector("startShapeData", startShapeData);
  print_shape_vector("endShapeData", endShapeData);
  print_shape_vector("weightShapeData", weightShapeData);

  auto blockDim = common_tiling_set_kernel_vec(tiling, context);
  COMMON_TILING_GET_PLATFORM;
  COMMON_TILING_GET_DTYPE_DATA(0);
  COMMON_TILING_ALIGN_DATA(256, size);
  // 分割数据
  COMMON_TILING_SPLIT(tiling);
  // 分割UB
  COMMON_TILING_GET_UB_SIZE;
  if (startShapeData[1] || endShapeData[1] || weightShape[1]) {
    ub_size /= 2;
  }
  // (start + end + weight + y) * double buffer = 4 * 2 = 8
  auto split_ub = ub_size / 8 / 256 * 256 / dtypeSize;
  if (dtype == ge::DataType::DT_FLOAT16) {
    // (start + end + weight + y + temp1 + temp2 + maxhalf + minHalf) * double
    // buffer =
    split_ub = ub_size / 16 / 256 * 256 / dtypeSize;
  }
  auto ret = splitUb(split_ub, formerLength, tailLength, alignDataSize,
                     alignedSize, size);
  tiling.set_formerTileLength(std::get<0>(ret));
  tiling.set_tailTileLength(std::get<1>(ret));
  tiling.set_formerFinalCalcCount(std::get<2>(ret));
  tiling.set_tailFinalCalcCount(std::get<3>(ret));
  tiling.set_finalKernelCaclCount(std::get<4>(ret));

  tiling_zero_workspace(context);

  tiling_save_to_buffer(tiling, context);

  print_common_tiling(tiling);

  return ge::GRAPH_SUCCESS;
}
} // namespace optiling

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext *context) {
  const gert::Shape *x1_shape = context->GetInputShape(0);
  gert::Shape *y_shape = context->GetOutputShape(0);
  *y_shape = *x1_shape;
  return GRAPH_SUCCESS;
}
} // namespace ge

namespace ops {
class Lerp : public OpDef {
public:
  explicit Lerp(const char *name) : OpDef(name) {
    this->Input("start")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Input("end")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Input("weight")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
    this->Output("y")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
        .Format({ge::FORMAT_ND, ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

    this->SetInferShape(ge::InferShape);

    this->AICore().SetTiling(optiling::TilingFunc);
    this->AICore().AddConfig("ascend310b");
  }
};

OP_ADD(Lerp);
} // namespace ops
